import pandas as pd
from logic.global_objs import g

class Backtest:
    def __init__(self, df):
        self.df = df
        #按日期做循环
        self.dates = self.df.index.unique()
        self.observers = []

    def onbar(self, index, date):
        df_bar = self.df.loc[date]
        if type(df_bar) is pd.Series:
            df_bar = df_bar.to_frame().T

        df_bar.index = df_bar['code']

        self.strategy.onbar(index, date, df_bar)


    #运行入口
    def run(self, s):
        '''
        1,每个日期是一个bar
        2,读出当前date的df_bar
        '''
        self.strategy = s
        g.notify(g.BKT_START,{})
        for index, date in enumerate(self.dates):
            self.onbar(index, date)
            g.notify(g.BKT_PROGRESS, {'rate': index/len(self.dates)})
        g.notify(g.BKT_FINISHED,{'name': s.name})

        return self.get_results()

    def get_results(self):
        rets = []
        date_start = None
        s = self.strategy
        df = s.acc.get_results_df()
        return df

if __name__ == '__main__':

    from codes.data_handler import DataHandler
    from engine.strategy import Strategy
    from engine.algos import RunOnce, RunPeriod, SelectFix, SelectBySignal, SelectTopK, WeightEqually, WeightFix

    MACD_EXP = '(EMA($close, 12) - EMA($close, 26))/$close - EMA((EMA($close, 12) - EMA($close, 26))/$close, 9)/$close'
    mom_5 = '$close/Ref($close, 5) - 1'
    rate = '$close/Ref($close, 1) - 1'

    features = [MACD_EXP, '$close', mom_5, rate]  # MACD
    names = ['MACD', '收盘价', '五日动量', 'rate']

    df = DataHandler().load('index_data', ['sh000300', 'sh000905', 'sz399006'], features=features, names=names)
    df = df['feature']
    df['to_buy'] = df['五日动量'] > 0
    df['to_sell'] = df['五日动量'] <= 0



    algo_list = [
        RunOnce(),
        SelectFix(instruments=['sh000300']),
        WeightEqually()
    ]

    algo_list2 = [
        RunOnce(),
        SelectFix(instruments=['sh000300', 'sh000905']),
        WeightFix(weights=[0.7, 0.3])
    ]

    algo_list3 = [
        RunPeriod(period=22),
        SelectFix(instruments=['sh000300', 'sh000905']),
        WeightFix(weights=[0.7, 0.3])
    ]

    algo_list_rolling = [
        SelectFix(instruments=['sh000300', 'sh000905', 'sz399006']),
        SelectBySignal(signal_buy='to_buy', signal_sell='to_sell'),
        SelectTopK(K=1,col='五日动量'),
        WeightEqually()
    ]

    s = Strategy(name='买入并持有_沪深300', algo_list=algo_list)
    s2 = Strategy(name='买入并持有_沪深300、中证500', algo_list=algo_list2)
    s3 = Strategy(name='买入并持有_沪深300、中证500_月再平衡', algo_list=algo_list3)
    s_rolling = Strategy(name='动量轮动', algo_list=algo_list_rolling)

    df = Backtest(df=df).run([s3, s_rolling])

    equity_cols = []
    for c in df.columns:
        df['equity_{}'.format(c)] = (df[c] + 1).cumprod()
        equity_cols.append('equity_{}'.format(c))

    df[equity_cols].plot()
    import matplotlib.pyplot as plt
    plt.rcParams['font.sans-serif'] = ['KaiTi']  # 指定默认字体
    plt.rcParams['axes.unicode_minus'] = False
    plt.show()

